# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com


import PricingEngineMod
from PricingEngineMod import *
reload(PricingEngineMod)

import ExplicitFiniteDifferencesEngineMod
from ExplicitFiniteDifferencesEngineMod import *
reload(ExplicitFiniteDifferencesEngineMod)

from math import sqrt, log

class BlackScholesExplicitFiniteDifferences(PricingEngine):

    def preComputations(self):
        p = dict(self.otherParams.p.items() + self.modelParams.p.items())
        strike = p['strike']
        maxTime = p['tenor']
        variables = []
        maxStock = 2*strike
        stock = FiniteDifferenceVariable('stock', 0, maxStock)
        variables.append(stock)
        time = FiniteDifferenceVariable('time', 0, maxTime)
        variables.append(time)

##        partitions_stock = 100
##        # need: dt < 0.5 (dS)^2
##        partitions_time = int(1.1 * 2 * maxTime / ((float(maxStock)/partitions_stock) ** 2))

        rate = p['rate']
        vol = p['vol']
        
        partitions_time = int(100)# * maxTime) #10 per year!
        partitions_stock = int(maxStock/2.0)#int(0.5 * 1 / (vol * sqrt(maxTime/partitions_time)))
        partitions = [partitions_stock, partitions_time]

        # PDE is V_t + rS V_S + 0.5 sig^2 S^2 V_SS - r V = 0
        #   V_t = 1/dt (V[i+1,j] - V[i,j])
        #   V_S = 1/(2dS) (V[i+1,j+1] - V[i+1,j-1])
        #   V_SS = 1/(dS)^2 (V[i+1,j+1] - 2 V[i+1,j] + V[i+1,j-1])
        #
        #   => V[i,j] = a_j V[i+1,j-1] + b_j V[i+1,j] + c_j V[i+1, j+1]
        #   where
        #       a_j = 1/(1+rdt) (0.5 sig^2 (x_j)^2 dt - 0.5 r x_j dt)
        #       b_j = 1/(1+rdt) (1 - sig^2 (x_j)^2 dt)
        #       c_j = 1/(1+rdt) (0.5 sig^2 (x_j)^2 dt + 0.5 r x_j dt)
        #   where
        #       x_j = S_j / dS
        
        
        dt = maxTime / partitions_time

        relations = []
        #1:        
        relation = Relation()
        relation.variableNames = ['stock', 'time']
        relation.numberSteps = [-1, 1]
        relation.func = lambda x, y : 1.0/(1+rate * x[1]) * 0.5 * (vol ** 2.0 * (y[0]/x[0]) ** 2.0 - rate * (y[0]/x[0])) * x[1]
        relations.append(relation)
        #2:        
        relation = Relation()
        relation.variableNames = ['stock', 'time']
        relation.numberSteps = [1, 1]
        relation.func = lambda x, y : 1.0/(1.0+rate * x[1]) * 0.5 * (vol ** 2.0 * (y[0]/x[0]) ** 2.0 + rate * (y[0]/x[0])) * x[1]
        relations.append(relation)
        #3:        
        relation = Relation()
        relation.variableNames = ['stock', 'time']
        relation.numberSteps = [0, 1]
        relation.func = lambda x, y : 1.0/(1.0+rate * x[1]) * (1.0 - vol ** 2.0 * (y[0]/x[0]) ** 2.0 * x[1]) 
        relations.append(relation)

        initials = {}
        initials['stock'] = p['initial']
        initials['time'] = 0

        self.fde = ExplicitFiniteDifferencesEngine(variables, partitions, 'time')
        # boundary conditions:
        #   t = Right: val = max(S-K,0)
        #   s = Left: val = 0
        #   s = Right: dval/ds = 1

        #1        
        b = BoundaryCondition()
        b.conditionNames = ['time']
        b.LeftOrRight = 'Right'
        if p['type'] == 'call':
            b.funct = lambda x : 0.0 if x[0] < p['strike'] else x[0] - p['strike']
        elif p['type'] == 'put':
            b.funct = lambda x : 0.0 if x[0] > p['strike'] else p['strike'] - x[0]
        else:
            raise Exception('unknown option type provided')
        self.fde.applyDirichletBoundaryCondition(b)

        #2
        b = BoundaryCondition()
        b.conditionNames = ['stock']
        b.LeftOrRight = 'Left'
        if p['type'] == 'call':
            b.funct = lambda x : 0.0
        elif p['type'] == 'put':
            b.funct = lambda x : p['strike']
        self.fde.applyDirichletBoundaryCondition(b)

        #3  V[i,j] = 1/(1+rdt) (V[i+1,j] + rS dt) at the upper limit if a call (i.e. dS = 1), 0 if a put
        b = BoundaryCondition()
        b.conditionNames = ['stock']
        b.LeftOrRight = 'Right'
        if p['type'] == 'call':
            vnrelations = []
            relation = Relation()
            relation.variableNames = ['time']
            relation.numberSteps = [1]
            relation.func = lambda x, y : 1.0/(1+rate * x[1])
            vnrelations.append(relation)
            relation = Relation()
            relation.variableNames = []
            relation.numberSteps = []
            relation.func = lambda x, y : 1.0/(1+rate * x[1]) * (y[0] * rate * x[1])
            vnrelations.append(relation)
            for relation in vnrelations:
                for varName in relation.variableNames:
                    relation.varIndices.append(self.fde.grid.findIndexInNamesOrder(varName))
            self.fde.applyVanNeumanBoundaryCondition(b, vnrelations)
        elif p['type'] == 'put':
            b.funct = lambda x : 0.0
            self.fde.applyDirichletBoundaryCondition(b)

        self.fde.initialise(relations, initials)

        
    def getPrice(self):
        if (self.currentStatus > 0):
            op = self.otherParams.p
            mp = self.modelParams.p
            return self.fde.doCalculation()